import os
import glob
from typing import List, Dict, Any
import json
import shutil
from datetime import datetime
import pandas as pd

def get_products_from_data_dir(data_dir: str) -> List[str]:
    """
    获取数据目录下的产品列表
    
    Args:
        data_dir: 数据目录
        
    Returns:
        产品名称列表
    """
    products = []
    for item in os.listdir(data_dir):
        item_path = os.path.join(data_dir, item)
        if os.path.isdir(item_path) and not item.startswith('.'):
            # 检查是否有必要的子目录结构
            train_dir = os.path.join(item_path, "train")
            test_dir = os.path.join(item_path, "test")
            if os.path.exists(train_dir) and os.path.exists(test_dir):
                products.append(item)
    
    return sorted(products)

def get_defect_types_from_train_dir(train_dir: str) -> List[str]:
    """
    获取训练目录中的缺陷类型列表
    
    Args:
        train_dir: 训练数据目录
        
    Returns:
        缺陷类型列表
    """
    defect_types = []
    for item in os.listdir(train_dir):
        item_path = os.path.join(train_dir, item)
        if os.path.isdir(item_path) and not item.startswith('.'):
            defect_types.append(item)
    
    return sorted(defect_types)

def count_images_in_dir(directory: str) -> int:
    """
    计算目录中的图片数量
    
    Args:
        directory: 目录路径
        
    Returns:
        图片数量
    """
    count = 0
    for ext in ['png', 'jpg', 'jpeg']:
        count += len(glob.glob(os.path.join(directory, f"*.{ext}")))
    
    return count

def get_product_stats(data_dir: str, product_name: str) -> Dict:
    """
    获取产品的统计信息
    
    Args:
        data_dir: 数据目录
        product_name: 产品名称
        
    Returns:
        统计信息字典
    """
    product_dir = os.path.join(data_dir, product_name)
    train_dir = os.path.join(product_dir, "train")
    test_dir = os.path.join(product_dir, "test")
    
    stats = {
        "name": product_name,
        "train": {},
        "test": {},
        "total_train_images": 0,
        "total_test_images": 0
    }
    
    # 统计训练集
    if os.path.exists(train_dir):
        defect_types = get_defect_types_from_train_dir(train_dir)
        for defect_type in defect_types:
            defect_dir = os.path.join(train_dir, defect_type)
            img_count = count_images_in_dir(defect_dir)
            stats["train"][defect_type] = img_count
            stats["total_train_images"] += img_count
    
    # 统计测试集
    if os.path.exists(test_dir):
        defect_types = get_defect_types_from_train_dir(test_dir)
        for defect_type in defect_types:
            defect_dir = os.path.join(test_dir, defect_type)
            img_count = count_images_in_dir(defect_dir)
            stats["test"][defect_type] = img_count
            stats["total_test_images"] += img_count
    
    return stats

def save_upload_file(uploaded_file, save_dir: str) -> str:
    """
    保存上传的文件
    
    Args:
        uploaded_file: 上传的文件对象
        save_dir: 保存目录
        
    Returns:
        保存的文件路径
    """
    os.makedirs(save_dir, exist_ok=True)
    file_path = os.path.join(save_dir, uploaded_file.name)
    
    with open(file_path, "wb") as f:
        f.write(uploaded_file.getbuffer())
    
    return file_path

def create_experiment_log(experiment_name: str, params: Dict, result_dir: str) -> str:
    """
    创建实验日志
    
    Args:
        experiment_name: 实验名称
        params: 实验参数
        result_dir: 结果目录
        
    Returns:
        日志文件路径
    """
    log_dir = os.path.join(result_dir, "logs")
    os.makedirs(log_dir, exist_ok=True)
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_file = os.path.join(log_dir, f"{experiment_name}_{timestamp}.json")
    
    log_data = {
        "experiment_name": experiment_name,
        "timestamp": timestamp,
        "params": params
    }
    
    with open(log_file, 'w', encoding='utf-8') as f:
        json.dump(log_data, f, ensure_ascii=False, indent=2)
    
    return log_file

def load_evaluation_results(result_dir: str) -> Dict:
    """
    加载评估结果
    
    Args:
        result_dir: 结果目录
        
    Returns:
        评估结果
    """
    summary_path = os.path.join(result_dir, "summary.json")
    if not os.path.exists(summary_path):
        return {}
    
    with open(summary_path, 'r', encoding='utf-8') as f:
        return json.load(f)

def get_experiment_logs(result_dir: str) -> List[Dict]:
    """
    获取实验日志列表
    
    Args:
        result_dir: 结果目录
        
    Returns:
        实验日志列表
    """
    logs = []
    log_dir = os.path.join(result_dir, "logs")
    
    if not os.path.exists(log_dir):
        return logs
    
    for log_file in os.listdir(log_dir):
        if log_file.endswith(".json"):
            log_path = os.path.join(log_dir, log_file)
            try:
                with open(log_path, 'r', encoding='utf-8') as f:
                    log_data = json.load(f)
                    log_data["file"] = log_file
                    logs.append(log_data)
            except:
                continue
    
    # 按时间戳排序
    logs.sort(key=lambda x: x.get("timestamp", ""), reverse=True)
    return logs 